import torch
from torch.utils.data import DataLoader

import models.generator as Gen
from data_set import MyDataSet
from models.detector_cnnlstm import CNNLSTMModel


def calculate_escape_rate(predictions, real_labels):
    """
    计算逃逸率
    Args:
        predictions: 形状为 [batch_size, 5] 的预测结果
        real_labels: 形状为 [batch_size, 5] 的真实标签
    Returns:
        fp: 假阳性数量
        total_pred_normal: 总的负样本数量
        escape_rate: 逃逸率
    """
    # 获取预测标签和真实标签的类别索引
    pred_classes = torch.argmax(predictions, dim=1)  # [batch_size]
    true_classes = torch.argmax(real_labels, dim=1)  # [batch_size]

    # 在这里，索引4代表正常流量
    # 假阳性(FP)：模型预测为正常但实际是恶意的样本数
    fp = ((pred_classes == 4) & (true_classes != 4)).sum().item()

    # 计算总的恶意样本数（负样本）
    total_malicious = (true_classes != 4).sum().item()

    # 计算逃逸率
    if total_malicious == 0:
        print("Warning: No malicious samples in batch")
        return 0, 0, 0

    escape_rate = fp / total_malicious

    # # 添加调试信息
    # print("Batch stats:")
    # print(f"- Total samples: {len(pred_classes)}")
    # print(f"- False positives: {fp}")
    # print(f"- Total malicious: {total_malicious}")
    # print(f"- Escape rate: {escape_rate:.4%}")

    return fp, total_malicious, escape_rate


# 计算并返回总体逃逸率
def escape_rate(model, test_dataloader, generator_model, ξ=0.1):
    """
    计算模型在真实数据和对抗样本上的逃逸率。
    model: 经过训练的检测模型（判别器）
    test_dataloader: 测试集DataLoader
    generator_model: 生成器模型
    ξ: 扰动约束
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    total_fp = 0
    total_pred_normal = 0
    escape_rates = []

    model.eval()  # 确保模型在评估模式下
    generator_model.eval()  # 确保生成器在评估模式下
    with torch.no_grad():  # 不计算梯度
        for batch_idx, (real_datas, real_labels) in enumerate(test_dataloader):
            real_datas = real_datas.to(device)  # 将数据转移到设备上
            real_labels = real_labels.to(device)  # 将标签转移到设备上
            # 生成对抗样本
            fake_samples = generator_model.generate_adv_samples(real_datas, real_labels, ξ=ξ)

            # 使用判别器对生成的对抗样本进行预测
            predictions = model(fake_samples)

            # 计算每个batch的逃逸率
            true_classes = torch.argmax(real_labels, dim=1)  # 获取真实标签的类别索引
            total_malicious = (true_classes != 4).sum().item()
            if total_malicious == 0:
                continue
            fp, total_normal, escape_rate_val = calculate_escape_rate(
                predictions, real_labels
            )
            total_fp += fp
            total_pred_normal += total_normal
            escape_rates.append(escape_rate_val)

        # 计算总体逃逸率
        if total_pred_normal == 0:  # 防止除以零
            avg_escape_rate = 0
        else:
            avg_escape_rate = total_fp / total_pred_normal

    #print(f"Total False Positive (FP): {total_fp}")
    #print(f"Total Predicted Normal: {total_pred_normal}")
    #print(f"Average Escape Rate: {avg_escape_rate:.4%}")  # 输出平均逃逸率
    return avg_escape_rate


# 主函数
def main():
    # 设置设备（CUDA或CPU）
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Using device: ", device)

    # 加载CNN-LSTM模型（判别器）
    trained_model = CNNLSTMModel().to(device)
    trained_model.load_state_dict(torch.load("./models/cnn_lstm_model.pth"))
    trained_model.eval()  # 设置为评估模式

    # 加载生成器模型
    g_input_dim = 118  # 输入生成器的维度
    g_output_dim = g_input_dim  # 生成器输出的维度
    feature_ranges = {i: (-float("inf"), float("inf")) for i in range(g_input_dim)}  # 特征取值范围
    
    generator_model = Gen.Generator(input_dim=g_input_dim, output_dim=g_output_dim, feature_ranges=feature_ranges).to(device)
    generator_model.load_state_dict(torch.load("./models/generator.pth"))
    generator_model.eval()  # 设置为评估模式

    # 加载数据集
    batch_size = 128
    test_data = MyDataSet('./data/train_subset.csv')
    test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=True)

    # 生成对抗样本并计算总体逃逸率
    escape_rate_val = escape_rate(model=trained_model, test_dataloader=test_dataloader, generator_model=generator_model, ξ=0.5)
    print(f"Overall Escape Rate: {escape_rate_val:.4%}")


if __name__ == "__main__":
    main()
